"""
JSON格式日志记录工具
将训练过程记录为结构化的JSON格式，便于分析
"""
import os
import json
from datetime import datetime


class JSONLogger:
    """
    JSON格式的训练日志记录器

    日志结构：
    {
        "metadata": {
            "model_name": "unet",
            "fold_idx": 0,
            "start_time": "2025-01-01 10:00:00",
            "config": {...}
        },
        "training": {
            "epochs": [
                {
                    "epoch": 0,
                    "train": {"loss": 0.5, "iou": 0.7, ...},
                    "val": {"loss": 0.4, "iou": 0.75, ...},
                    "lr": 0.001,
                    "time": 120.5
                },
                ...
            ]
        },
        "best_model": {
            "epoch": 45,
            "metrics": {...},
            "checkpoint_path": "..."
        },
        "summary": {
            "total_epochs": 100,
            "total_time": 12000,
            "best_miou": 0.85,
            "end_time": "2025-01-01 12:00:00"
        }
    }
    """

    def __init__(self, log_dir, model_name, fold_idx=None):
        """
        初始化JSON日志记录器

        Args:
            log_dir: 日志保存目录
            model_name: 模型名称
            fold_idx: 折数索引
        """
        self.log_dir = log_dir
        self.model_name = model_name
        self.fold_idx = fold_idx

        # 创建日志目录
        os.makedirs(log_dir, exist_ok=True)

        # 构建日志文件名
        if fold_idx is not None:
            self.log_filename = f"{model_name}_fold_{fold_idx}.json"
        else:
            self.log_filename = f"{model_name}.json"

        self.log_path = os.path.join(log_dir, self.log_filename)

        # 初始化日志数据结构
        self.log_data = {
            "metadata": {
                "model_name": model_name,
                "fold_idx": fold_idx,
                "start_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "config": {}
            },
            "training": {
                "epochs": []
            },
            "best_model": {
                "epoch": -1,
                "metrics": {},
                "checkpoint_path": ""
            },
            "summary": {}
        }

        print(f"[INFO] JSON日志初始化完成: {self.log_path}")

    def log_config(self, config):
        """
        记录配置信息

        Args:
            config: 配置字典
        """
        self.log_data["metadata"]["config"] = config
        self._save()
        print(f"[INFO] 配置已记录到JSON日志")

    def log_epoch(self, epoch, train_metrics, val_metrics, lr, epoch_time):
        """
        记录单个epoch的训练信息

        Args:
            epoch: epoch编号
            train_metrics: 训练集指标字典
            val_metrics: 验证集指标字典
            lr: 当前学习率
            epoch_time: epoch用时（秒）
        """
        epoch_data = {
            "epoch": epoch,
            "train": train_metrics,
            "val": val_metrics,
            "lr": lr,
            "time": epoch_time,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }

        self.log_data["training"]["epochs"].append(epoch_data)
        self._save()

    def log_best_model(self, epoch, metrics, checkpoint_path):
        """
        记录最佳模型信息

        Args:
            epoch: 最佳模型的epoch
            metrics: 最佳模型的指标
            checkpoint_path: 检查点路径
        """
        self.log_data["best_model"] = {
            "epoch": epoch,
            "metrics": metrics,
            "checkpoint_path": checkpoint_path,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        self._save()
        print(f"[INFO] 最佳模型信息已更新: Epoch {epoch}, mIoU={metrics.get('miou', 0):.4f}")

    def log_summary(self, total_epochs, total_time, best_miou):
        """
        记录训练总结信息

        Args:
            total_epochs: 总epoch数
            total_time: 总训练时间（秒）
            best_miou: 最佳mIoU
        """
        self.log_data["summary"] = {
            "total_epochs": total_epochs,
            "total_time": total_time,
            "best_miou": best_miou,
            "end_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        self._save()
        print(f"[INFO] 训练总结已记录")

    def _save(self):
        """保存日志到文件"""
        with open(self.log_path, 'w', encoding='utf-8') as f:
            json.dump(self.log_data, f, indent=2, ensure_ascii=False)

    def get_log_path(self):
        """获取日志文件路径"""
        return self.log_path


def create_json_logger(log_dir, model_name, fold_idx=None):
    """
    创建JSON日志记录器（工厂函数）

    Args:
        log_dir: 日志目录
        model_name: 模型名称
        fold_idx: 折数索引

    Returns:
        JSONLogger实例
    """
    return JSONLogger(log_dir, model_name, fold_idx)
